# utf-8

import tensorflow as tf
import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
import os
from tensorflow import keras
import sys
import cv2

#展示版本信息
# print(tf.__version__)
# print(sys.version_info)
# for module in mpl, pd, np, keras:
#     print(module.__name__, module.__version__)

# 设置图片路径
train_dir = 'D:/python_exercise/tensorflowExperiment/catAndDog/input/train/'
valid_dir = 'D:/python_exercise/tensorflowExperiment/catAndDog/input/validation/'
test_dir = 'D:/python_exercise/tensorflowExperiment/catAndDog/input/test1/'

# 设置参数
height = 128
width = 128
channel = 3#图片通道
batch_size = 64#批量大小
valid_batch_size = 64#验证批量大小
num_classes = 2#训练类别
epochs = 2#迭代次数


# 训练模型
def trainModel(model, train_generator, valid_generator, callbacks):
    history = model.fit(
        train_generator,
        epochs=epochs,
        validation_data = valid_generator,
        callbacks = callbacks
    )
    return history

# 显示训练过程中损失和准确性的变化
def plot_learning_curves(history, label, epochs, min_value, max_value):
    data = {}
    data[label] = history.history[label]
    data['val_' + label] = history.history['val_' + label]
    pd.DataFrame(data).plot(figsize=(8, 5))
    plt.grid(True)
    plt.axis([0, epochs, min_value, max_value])
    plt.show()


# 完成训练后使用该模型对图像进行分类，并将其保存到两个不同的文件夹中
def predictModel(model, output_model_file):

    # 加载模型的权重
    model.load_weights(output_model_file)

    os.makedirs('./save', exist_ok=True)
    os.makedirs('./save/cat', exist_ok=True)
    os.makedirs('./save/dog', exist_ok=True)

    test_dir = 'D:/python_exercise/tensorflowExperiment/catAndDog/input/test1/'  # 1-12500.jpg
    for i in range(1, 12500):
        img_name = test_dir + '{}.jpg'.format(i)
        img = cv2.imread(img_name)
        img = cv2.resize(img, (width, height))
        img_arr = img / 255.0
        img_arr = img_arr.reshape((1, width, height, 3))
        pre = model.predict(img_arr)
        #pre为识别结果数组，分别表示该图片属于猫和狗的概率。
        #最后根据pre数组中第一个元素和第二个元素的大小关系来判断该图片属于猫还是狗。
        #如果pre[0][0] > pre[0][1]，则认为该图片属于猫；否则认为该图片属于狗 
        if pre[0][0] > pre[0][1]:
            cv2.imwrite('./save/cat/' + '{}.jpg'.format(i), img)
            print(img_name, ' 属于猫！')
        else:
            cv2.imwrite('./save/dog/' + '{}.jpg'.format(i), img)
            print(img_name, ' 属于狗')


if __name__ ==  '__main__':

    print('开始导入数据...')

    # 导入数据并进行扩充
    train_datagen = keras.preprocessing.image.ImageDataGenerator(
        rescale=1. / 255,
        rotation_range=40,
        width_shift_range=0.2,
        height_shift_range=0.2,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True,
        fill_mode='nearest',
    )

    train_generator = train_datagen.flow_from_directory(
        train_dir,
        target_size=(width, height),
        batch_size=batch_size,
        seed=7,
        shuffle=True,
        class_mode='categorical'
    )

    valid_datagen = keras.preprocessing.image.ImageDataGenerator(
        rescale=1. / 255,
    )

    valid_generator = valid_datagen.flow_from_directory(
        valid_dir,
        target_size=(width, height),
        batch_size=valid_batch_size,
        seed=7,
        shuffle=False,
        class_mode="categorical"
    )

    #训练数据量
    train_num = train_generator.samples
    valid_num = valid_generator.samples
    print(train_num, valid_num)

    print('建立模型...')

    # 建立模型
    model = keras.models.Sequential([
        keras.layers.Conv2D(filters=32, kernel_size=3,
                            padding='same', activation='relu',
                            input_shape=[width, height, channel]),
        keras.layers.Conv2D(filters=32, kernel_size=3,
                            padding='same', activation='relu'),
        keras.layers.MaxPool2D(pool_size=2),

        keras.layers.Conv2D(filters=64, kernel_size=3,
                            padding='same', activation='relu'),
        keras.layers.Conv2D(filters=64, kernel_size=3,
                            padding='same', activation='relu'),
        keras.layers.MaxPool2D(pool_size=2),

        keras.layers.Conv2D(filters=128, kernel_size=3,
                            padding='same', activation='relu'),
        keras.layers.Conv2D(filters=128, kernel_size=3,
                            padding='same', activation='relu'),
        keras.layers.MaxPool2D(pool_size=2),

        keras.layers.Flatten(),
        keras.layers.Dense(128, activation='relu'),
        keras.layers.Dense(num_classes, activation='softmax')
    ])

    model.compile(loss='categorical_crossentropy',
                  optimizer='adam', metrics=['accuracy'])
    model.summary()

    # 设置模型保存路径
    logdir = './graph_def_and_weights'
    if not os.path.exists(logdir):
        os.mkdir(logdir)

    output_model_file = os.path.join(logdir,"catDog_weights.h5")

    print('开始训练 ...')
    

    mode = input('选择模式: 1.训练 2.预测\n输入选项编号: ')
    if mode == '1':
        callbacks = [
            keras.callbacks.TensorBoard(logdir),
            keras.callbacks.ModelCheckpoint(output_model_file,
                                            save_best_only=True,
                                            save_weights_only=True),
            keras.callbacks.EarlyStopping(patience=5, min_delta=1e-3)
        ]
        history = trainModel(model, train_generator, valid_generator, callbacks)
        plot_learning_curves(history, 'accuracy', epochs, 0, 1)
        plot_learning_curves(history, 'loss', epochs, 0, 5)
    elif mode == '2':
        # 此模式仅供已有训练好的模型时使用.
        predictModel(model, output_model_file)
    else:
        print('请输入正确的选项.')

    print('完成! Exit.')
